/*
 * Simple Java Encryptor
 * 
 * File: SimpleJavaEncryptor.java
 * 
 * Date: 28/07/2013
 * 
 * @author: Aarón C.d.C (AKA Magnus or Ninjihaku)
 * 
 * Description: Java application that can be used to encrypt files.
 * 
 */
package simpleencryptor;

import java.io.*;
import java.security.SecureRandom;
import java.io.File;

/* Class used to handle FileStreams */
class FileStreams {
    public static RandomAccessFile FI; //Input
    public static RandomAccessFile FO; //Output
    //Stream locks
    public static boolean IsInuse = false;
    public static boolean IsReady = false;
    //The file checksum is also stored here.
    public static byte[] ChkSum = new byte[2];
    
    /* Function that start the file streams */
    public static void StartStreams(String fileInput, String fileOutput)
    {
        try{
        FI = new RandomAccessFile(fileInput, "rw");
        FO = new RandomAccessFile(fileOutput, "rw");
        }catch(Exception e){return;}
        
        IsReady = true;
    }
    
}

/* Class that contain the Decryption routines */
class DecryptClass extends Thread {
    int cores, processed = 0;
    long memory;
    double bufSize;
    byte chunks;
    double chunkIndex;
    byte[] keys, Chk = new byte[2], ChkOr = new byte[2];
    char[] password;
    
    /* Function that calculate the number of chunks to decrypt */
    public double GetNumberOfChunks(String FilePath, long bufSize)
    {
        RandomAccessFile TAF;
        try{
         TAF = new RandomAccessFile(FilePath, "r");
         double res = Math.ceil(TAF.length() / bufSize);
         TAF.close(); TAF = null;
         return res;
        }catch (Exception e){ return 0; }
    }
    
    /* Function that alocate a portion of the memory to be used by a thread */
    public byte[] AlocateThreadMemory()
    {
        try{
        byte[] mem = new byte[(int)bufSize];
        return mem;
        }catch (Exception e){ return null; }
    }
    
    /* Class constructor */
    DecryptClass(int cores_,long mem_, double bSize_, double chunkIndex_, byte[] keys_, char[] password_, byte[] chksm)
    {
        cores = cores_;
        memory = mem_;
        bufSize = bSize_;
        chunkIndex = chunkIndex_;
        keys = keys_;
        password = password_;
        ChkOr = chksm;
    }
    
    /* The whole decryption process */
    @Override
    public void run()
    {
        try{
        long fsize = FileStreams.FI.length() - 130;
        int cnt = (int)(bufSize * chunkIndex)%password.length, cnt2 = (int)(bufSize * chunkIndex)%keys.length;
        byte[] data = AlocateThreadMemory();
        
        //If the memory is bigger than the chunk, realocate the memory to match the chunk size
        if((bufSize * chunkIndex)+data.length > fsize){
            data = new byte[(int)(fsize - (bufSize * chunkIndex))];
            System.out.println("Thread " + (int)chunkIndex + " realocated with " + 
                    (int)((fsize - (bufSize * chunkIndex))/1024) + " kB of memory.");
        }
        
        //Wait until the file streams are ready to use
        while(FileStreams.IsInuse){ }
        
        FileStreams.IsInuse = true; //Lock streams
        FileStreams.FI.seek((long)(bufSize * chunkIndex));
        FileStreams.FI.read(data);
        FileStreams.IsInuse = false; //Unlock streams
        
        
        System.out.println("Thread " + (int)chunkIndex + " seeks for pos 0x" + Integer.toHexString((int)(bufSize * chunkIndex)));
        
        //The whole decryption process
        for(int x = 0; x < data.length; x++)
        {
            data[x] = (byte)((keys[cnt2] ^ password[cnt] >> 1) ^ (password[cnt] << 1) ^ data[x]);
            Chk[(cnt % 2 == 0)?0:1] += data[x];
            cnt = (cnt == password.length-1)?0:cnt + 1;
            cnt2 = (cnt2 == keys.length-1)?0:cnt2 + 1;
        }
        
        //Wait for this thread's turn to write the information
        
        if(FileStreams.FO.length() < (bufSize * chunkIndex))
            System.out.println("THREAD " + (int)chunkIndex + " IS WAITING...");
        while(FileStreams.FO.length() != (bufSize * chunkIndex)){
            if(FileStreams.FO.length() > (bufSize * chunkIndex))
            {
                //If the address has been written, then there is a big problem!
                System.out.println("FATAL ERROR IN THREAD " + (int)chunkIndex + " - SOMEONE WROTE IN MY ADDRESS!");
                Thread.currentThread().interrupt();
            }
        }
        
        //Wait for the streams to be unlocked. Then, write the information of the chunk.
        while(FileStreams.IsInuse){}
        System.out.println((int)chunkIndex + " Storing data at 0x" + Integer.toHexString((int)(bufSize * chunkIndex)));
        FileStreams.IsInuse = true;
        FileStreams.FO.seek((long)(bufSize * chunkIndex));
        FileStreams.FO.write(data);
            FileStreams.ChkSum[0] += Chk[0]; FileStreams.ChkSum[1] += Chk[1];
        FileStreams.IsInuse = false;
        Thread.currentThread().interrupt();
        }catch(Exception e){ System.out.println(e); }
    }
    
    
}

/* Class that contain the Encryption routines */
class EncryptClass extends Thread {
    int cores, processed = 0;
    long memory;
    double bufSize;
    byte chunks;
    double chunkIndex;
    byte[] keys, Chk = new byte[2];
    char[] password;
    
    public double GetNumberOfChunks(String FilePath, long bufSize)
    {
        RandomAccessFile TAF;
        try{
         TAF = new RandomAccessFile(FilePath, "r");
         double res = Math.ceil(TAF.length() / bufSize);
         TAF.close();
         return res;
        }catch (Exception e){ return 0; }
    }
    
    public byte[] AlocateThreadMemory()
    {
        try{
        byte[] mem = new byte[(int)bufSize];
        return mem;
        }catch (Exception e){ return null; }
    }
    
    EncryptClass(int cores_,long mem_, double bSize_, double chunkIndex_, byte[] keys_, char[] password_)
    {
        cores = cores_;
        memory = mem_;
        bufSize = bSize_;
        chunkIndex = chunkIndex_;
        keys = keys_;
        password = password_;
    }
    
    @Override
    public void run()
    {
        try{
        int cnt = (int)(bufSize * chunkIndex)%password.length, cnt2 = (int)(bufSize * chunkIndex)%keys.length;
        byte[] data = AlocateThreadMemory();
        
        if((bufSize * chunkIndex)+data.length > FileStreams.FI.length()){
            data = new byte[(int)(FileStreams.FI.length() - (bufSize * chunkIndex))];
            System.out.println("Thread " + (int)chunkIndex + " realocated with " + 
                    (int)((FileStreams.FI.length() - (bufSize * chunkIndex))/1024) + " kB of memory.");
        }
        
        while(FileStreams.IsInuse){ }
        
        FileStreams.IsInuse = true;
        FileStreams.FI.seek((long)(bufSize * chunkIndex));
        FileStreams.FI.read(data);
        FileStreams.IsInuse = false;
        
        
        System.out.println("Thread " + (int)chunkIndex + " seeks for pos 0x" + Integer.toHexString((int)(bufSize * chunkIndex)));
        
        for(int x = 0; x < data.length; x++)
        {
            Chk[(cnt % 2 == 0)?0:1] += data[x];
            data[x] = (byte)(data[x] ^ (password[cnt] << 1) ^ (keys[cnt2] ^ password[cnt] >> 1));
            cnt = (cnt == password.length-1)?0:cnt + 1;
            cnt2 = (cnt2 == keys.length-1)?0:cnt2 + 1;
        }
        if(FileStreams.FO.length() < (bufSize * chunkIndex))
            System.out.println("THREAD " + (int)chunkIndex + " IS WAITING...");
        while(FileStreams.FO.length() != (bufSize * chunkIndex)){
            if(FileStreams.FO.length() > (bufSize * chunkIndex))
            {
                System.out.println("FATAL ERROR IN THREAD " + (int)chunkIndex + " - SOMEONE WROTE IN MY ADDRESS!");
                Thread.currentThread().interrupt();
            }
        }
        
        while(FileStreams.IsInuse){}
        System.out.println((int)chunkIndex + " Storing data at 0x" + Integer.toHexString((int)(bufSize * chunkIndex)));
        FileStreams.IsInuse = true;
        FileStreams.FO.seek((long)(bufSize * chunkIndex));
        FileStreams.FO.write(data);
            FileStreams.ChkSum[0] += Chk[0]; FileStreams.ChkSum[1] += Chk[1];
        FileStreams.IsInuse = false;
        Thread.currentThread().interrupt();
        }catch(Exception e){ System.out.println(e); }
    }
    
    
}

/* Main class */
public class SimpleEncryptor {

    /**
     * @param args the command line arguments
     */
    static int cores, processed = 0;
    static long memory;
    static double bufSize;
    static byte chunks;
    static RandomAccessFile FI;
    static RandomAccessFile FO;
    static boolean isInuse = false;
    static byte[] ChkSm = new byte[2];
    static byte MemReserve = 40;
    
    /* Routine that generate a 128 bit encryption secure key */
    public static byte[] GenerateKey() throws Exception{
        byte[] Keys = new byte[128];
        SecureRandom RNG = SecureRandom.getInstance("SHA1PRNG");
        RNG.nextBytes(Keys);
       return Keys;
    }
    
    /* Routine that print the help in the screen */
    public static void PrintHelp()
    {
        System.out.println("\n\nSimple Encryptor"
                + "\n\n"
                + " Arguments: [if] [password] [of] [mode] (reserve)"
                + "\n\n"
                + "     [if]: Path of the input file to be processed.\n\n"
                + "     [password]: A valid password. Must be at least 8 character long, and at most 128 character long.\n\n"
                + "     [of]: Path of the output file to be generated.\n\n"
                + "     [mode]: Must be 0 or ENCRYPT for encryption, or 1 or DECRYPT for decryption."
                + "     (reserve): A percentage of memory to reserve for each thread. Must be a value between 20 and 100."
                + "     By default or if ommited, it takes a 40%."
                + "\n\n"
                + "When decrypting, the file must be a file that was also encrypted by this program, and the given"
                + " password must be the same as the one provided during the encryption process. Otherwishe, it"
                + " will fail and the decrypted file will be wrong.\n\n");
    }
    
    public static double GetNumberOfChunks(double bufSize)
    {
        try{ 
         double res = Math.ceil(FileStreams.FI.length() / bufSize);
         return res;
        }catch (Exception e){ return 0; }
    }
    
    public static void main(String[] args) throws IOException {
        boolean d = false; //d is a bitflag used during the encryption process.
        
        //Get some system information, like the cores.
        cores = Runtime.getRuntime().availableProcessors();
        
        //Check if every argument is correct.
        System.out.println("Detected " + cores + " cores.");
        if(args.length == 0)
        {
            PrintHelp();
            return;
        }
        if(((args.length > 1 && args.length < 4)||args.length > 5) && args.length != 0)
        {
            System.out.println("FATAL ERROR: Missing arguments!"); return;
        }
        if(args[0].equals("/?")){
            PrintHelp();
            return;
        }
        if(args[0].equals("") || args[0] == null)
        {
            System.out.println("FATAL ERROR: No input file specified!"); return;
        }
        if(args[1].equals("") || args[1] == null)
        {
            System.out.println("FATAL ERROR: No password specified!");return;
        }
        if(args[1].length() < 8)
        {
            System.out.println("ERROR: The password must be at least 8 characters!");return;
        }
        if(args[1].length() > 128)
        {
            System.out.println("ERROR: The password must be at most 128 characters!");return;
        }
        if(args[2].equals("") || args[2] == null)
        {
            System.out.println("FATAL ERROR: No output file specified!");return;
        }
        if(args[3].equals("") || args[3] == null)
        {
            System.out.println("FATAL ERROR: No mode specified!");return;
        }
        if(args.length == 5)
            if(Integer.valueOf(args[4]) <= 100 && Integer.valueOf(args[4]) >= 20)
                MemReserve = Byte.valueOf(args[4]);
        try{
            
       //Get the memory information of the VM
            
        memory = Runtime.getRuntime().totalMemory();
        System.out.println("Avalible memory: " + (int)((memory/1024)/1024) + " MB");
        System.out.println("Reserved memory: " + MemReserve + "%");
        FileStreams.StartStreams(args[0], args[2]);
        bufSize = (Math.ceil(memory/cores)-(Math.ceil(memory/cores))*(0.01*MemReserve));
        
        chunks = (byte) GetNumberOfChunks(bufSize);
        }catch(Exception e){return;}
        
        //The encryption routines
        if(args[3].toLowerCase().equals("encrypt") || args[3].equals("0") && FileStreams.IsReady)
        {
            
            //If this bitflag is false at the end, it means that the encryption process
            //was unsuccessfull. For some reason it fails sometimes, so the program ensures
            //that the encryption was successfull, and it it wasn't, it starts again from
            //the beggining.
            
            boolean intok = false;
            
            //The loop.
            while(intok == false){
                intok = false;
                processed = 0;
                
                //Initialize the streams.
                FileStreams.FI.seek(0);
                FileStreams.FO.seek(0);
            
            System.out.println("Starting...");
            
            try{
                
            byte[] ks = GenerateKey(); //Generate a new key
            
            Thread[] Threading = new Thread[cores]; //1 thread per core.
            
            System.out.println("Assigning processes...");
            
            for(int x = 0; x < cores && x < chunks; x++)
            {
                
                //Just a very small wait to ensure everything synch.
                Thread.sleep(25);
                
                //Create, assign and start each thread
                EncryptClass l = new EncryptClass(cores, memory, bufSize, processed, ks, args[1].toCharArray());
                Threading[x] = new Thread(l);
                Threading[x].start();
                
                System.out.println("Process " + (x+1) + " assigned of " + cores + " processes with " + (int)(bufSize/1024) + " kB of memory.");
                processed++;
            }
            
            //Look for free threads, and assign a new one to them.
            while(processed < chunks){
                for(int x = 0; x < cores && x < chunks; x++)
                {
                    Thread.sleep(25);
                    if(Threading[x].getState() == Thread.State.TERMINATED)
                    {
                        System.out.println("Process " + (x+1) + " allocated " + (int)(bufSize/1024) + " kB of memory for " + processed);
                        EncryptClass l = new EncryptClass(cores, memory, bufSize, processed, ks, args[1].toCharArray());
                        Threading[x] = new Thread(l);
                        Threading[x].start(); 
                        processed++;
                    }
                }
            }
            
            //Wait until the length of the file is the same as the original one.
            
            while(FileStreams.FO.length() != FileStreams.FI.length()){ }
            
            System.out.println("Written " + (FileStreams.FI.length()/1024) + " kB of data");
            
            //Write the checksum and the key information.
            
            while(FileStreams.IsInuse){}
            System.out.println("Checksum: 0x" + Integer.toHexString(FileStreams.ChkSum[0]) + Integer.toHexString(FileStreams.ChkSum[1]));
            FileStreams.IsInuse = true;
            FileStreams.FO.write(FileStreams.ChkSum);
            FileStreams.FO.write(ks);
            FileStreams.IsInuse = false;
            
            //Check again for a file integrity. If it is not ok, it restart the process.
            
            if(FileStreams.FI.length() == FileStreams.FO.length()-130){
                System.out.println("Integrity check OK");
                intok = true;
            }else{
                System.out.println("INTEGRITY CHECK FAILED! RETRYING... ");
                FileStreams.FO.close();
                Thread.sleep(25);
                File x = new File(args[2]); x.setWritable(true); x.delete();
                Thread.sleep(25);
                FileStreams.StartStreams(args[0], args[2]);
                FileStreams.ChkSum[0] = 0; FileStreams.ChkSum[1] = 0;
            }
            
            }catch(Exception e){ System.out.println(e); }
            }
            
            //Everything is fine. Finish.
            System.out.println("\n~- FINISH -~");
            
        }else if(args[3].toLowerCase().equals("decrypt") || args[3].equals("1") && FileStreams.IsReady)
        {
            
            //Now, the decryption process
            
            System.out.println("Starting...");
            try{
            byte[] ks = new byte[128];
            
            //Retrieve the checksum and the keys stored in a previous encryption
            
            FileStreams.IsInuse = true;
            FileStreams.FI.seek(FileStreams.FI.length()-130);
            FileStreams.FI.read(ChkSm);
            FileStreams.FI.read(ks);
            FileStreams.FI.seek(0);
            FileStreams.IsInuse = false;    
               
            //And now, everything is pretty much the same. The multi-core processing.
            
            Thread[] Threading = new Thread[cores];
            System.out.println("Assigning processes...");
            for(int x = 0; x < cores && x < chunks; x++)
            {
                //Thread.sleep(25);
                DecryptClass l = new DecryptClass(cores, memory, bufSize, processed, ks, args[1].toCharArray(),ChkSm);
                Threading[x] = new Thread(l);
                Threading[x].start();
                System.out.println("Process " + (x+1) + " assigned of " + cores + " processes with " + (int)(bufSize/1024) + " kB of memory.");
                processed++;
            }
            while(processed < chunks){
                for(int x = 0; x < cores && x < chunks; x++)
                {
                    //System.out.flush();
                    if(Threading[x].getState() == Thread.State.TERMINATED)
                    {
                        System.out.println("Process " + (x+1) + " allocated " + (int)(bufSize/1024) + " kB of memory for " + processed);
                        DecryptClass l = new DecryptClass(cores, memory, bufSize, processed, ks, args[1].toCharArray(),ChkSm);
                        Threading[x] = new Thread(l);
                        Threading[x].start(); 
                        processed++;
                    }
                }
            }
            
            while(FileStreams.FO.length() != FileStreams.FI.length()-130){ }
            
            System.out.println("Written " + ((FileStreams.FI.length()-130)/1024) + " kB of data");
            
            }catch(Exception e){ System.out.println(e); }
            
            // Verify the checksum and inform the user about the result.
            
            if(ChkSm[0] == FileStreams.ChkSum[0] && ChkSm[1] == FileStreams.ChkSum[1])
            System.out.println("Checksum: 0x" + Integer.toHexString(FileStreams.ChkSum[0]) + Integer.toHexString(FileStreams.ChkSum[1])
                    + " verification is CORRECT");
            else
                System.out.println("CHECKSUM VERFICATION FAILED!");
                
            //We're finish again
            
            System.out.println("\n~- FINISH -~");
            
        }else{
            
            //This happens if something is wrong with the parameter specified
            
            System.out.println("ERROR: The operation is invalid, or the file is inaccessible!");
        }
    }
}
